local math = require "math"

local bst = require "bfmoon.bst"

local avl = bst:new()

local function right_rotation(root)
	local parent = root.parent
	local pivot = root.left
	if root.left.right then
		root.left = root.left.right
		root.left.parent = root
	end
	pivot.right = root
	root.parent = pivot
	if parent then
		if parent.left == root then
			parent.left = pivot
		else
			parent.right = pivot
		end
		pivot.parent = parent
	end
	return pivot
end

local function left_rotation(root)
	parent = root.parent
	pivot = root.right
	if root.right.left then
		root.right = root.right.left
		root.right.parent = root
	end
	pivot.left = root
	root.parent = pivot
	if parent.left == root then
		parent.left = pivot
	else
		parent.right = pivot
	end
	pivot.parent = parent
	return pivot
end

local function balance_node(node)
	if node.balance == -2 then
		if node.right.balance == 1 then
			node = right_rotation(node.right)
		end
		left_rotation(node)
	else
		if node.left.balance == -1 then
			node = left_rotation(node.left)
		end
		right_rotation(node)
	end
end

function avl:insert(value, replace)
	local replaced, node = bst.insert(self, value, replace)
	if not replaced then
		node.balance = 0
		while node.parent do
			if node.parent.left == node then
				node.parent.balance = node.parent.balance + 1
			else
				node.parent.balance = node.parent.balance - 1
			end
			node = node.parent
			if math.abs(node.balance) > 1 then
				balance_node(node)
			end
		end
	end
end

return avl
